# Databricks notebook source
# MAGIC %md
# MAGIC # Installations

# COMMAND ----------

# MAGIC %pip install -U mostlyai[local] 
# MAGIC %pip install ipywidgets
# MAGIC %restart_python

# COMMAND ----------

import pandas as pd
from mostlyai.sdk import MostlyAI

mostly = MostlyAI(local=True)

# COMMAND ----------

# MAGIC %md
# MAGIC # Read and process tables

# COMMAND ----------

# MAGIC %md
# MAGIC ## Personas

# COMMAND ----------

#ALL ROWS DISTINCT
df_personas = spark.table("admin_govern_sta_des.tmp_mostlyai.personas")
display(df_personas.limit(10))

# COMMAND ----------

# MAGIC %md
# MAGIC ## Documents Personals

# COMMAND ----------

#ALL ROWS DISTINCT
#SAME ROWID_OBJECT_PERSONA MULTIPLE IDENTIFICADOR
df_documents_personals = spark.table("admin_govern_sta_des.tmp_mostlyai.documents_personals")
display(df_documents_personals.limit(10))

# COMMAND ----------

# MAGIC %md
# MAGIC ## Persona Adreça

# COMMAND ----------

#MULTIPLE ROWID_OBJECT_PERSONA W/ MULTIPLE IDENTIFICADOR
df_persona_adreca = spark.table("admin_govern_sta_des.tmp_mostlyai.`persona-adreca`")
display(df_persona_adreca.limit(10))

# COMMAND ----------

from pyspark.sql.functions import lit
df_persona_adreca = df_persona_adreca.withColumn("const", lit("A"))
display(df_persona_adreca)

# COMMAND ----------

# MAGIC %md
# MAGIC ## Adreça

# COMMAND ----------

df_adreca = spark.table("admin_govern_sta_des.tmp_mostlyai.adreca")
display(df_adreca.limit(10))

# COMMAND ----------

# MAGIC %md
# MAGIC ## Contactes

# COMMAND ----------

#ROWID_OBJECT_PERSONA W/ MULTIPLE IDENTIFICADOR (SOMETIMES MORE THAN 100, ONCE W/ 614"!!)
df_contactes = spark.table("admin_govern_sta_des.tmp_mostlyai.contactes")
display(df_contactes.limit(10))

# COMMAND ----------

# MAGIC %md
# MAGIC # Generator

# COMMAND ----------

# MAGIC %md
# MAGIC ## Configuration

# COMMAND ----------

from mostlyai.sdk.domain import ModelEncodingType

personas_table_config = {
    "name": "personas",
    "data": df_personas,
    "tabular_model_configuration": {
        "max_training_time": 45
    },
    "primary_key": "ROWID_OBJECT_PERSONA",
    "columns": [{"name" : "ROWID_OBJECT_PERSONA", "model_encoding_type": ModelEncodingType.auto},
                {"name" : "NOM", "model_encoding_type": ModelEncodingType.tabular_character},
                {"name" : "COGNOM1", "model_encoding_type": ModelEncodingType.tabular_character},
                {"name" : "COGNOM2", "model_encoding_type": ModelEncodingType.tabular_character}]
}

documents_personals_table_config = {
    "name": "documents_personals",
    "data": df_documents_personals,
    "tabular_model_configuration": {
        "max_training_time": 45
    },
    "foreign_keys": [{"column": "ROWID_OBJECT_PERSONA", "referenced_table": "personas", "is_context": True}],
    "columns": [{"name" : "ROWID_OBJECT_PERSONA", "model_encoding_type": ModelEncodingType.auto},
                {"name" : "IDENTIFICADOR", "model_encoding_type": ModelEncodingType.tabular_character}]
}

# contactes_config = {
#     "name": "contactes",
#     "data": df_contactes,
#     "tabular_model_configuration": {
#         "max_training_time": 10
#     },
#     "foreign_keys": [{"column": "ROWID_OBJECT_PERSONA", "referenced_table": "personas", "is_context": True}],
#     "columns": [{"name" : "ROWID_OBJECT_PERSONA", "model_encoding_type": ModelEncodingType.auto},
#                 {"name" : "IDENTIFICADOR", "model_encoding_type": ModelEncodingType.tabular_character}]
# }


persona_adreca_config = {
    "name": "persona_adreca",
    "data": df_persona_adreca,
    "tabular_model_configuration": {
        "max_training_time": 45
    },
    "foreign_keys": [{"column": "ROWID_OBJECT_PERSONA", "referenced_table": "personas", "is_context": True},
                     {"column": "ROWID_OBJECT_ADRECA", "referenced_table": "adreca", "is_context": False}],
    "columns": [{"name" : "ROWID_OBJECT_PERSONA", "model_encoding_type": ModelEncodingType.auto},
                {"name" : "ROWID_OBJECT_ADRECA", "model_encoding_type": ModelEncodingType.auto}]
                #{"name" : "const", "model_encoding_type": ModelEncodingType.tabular_categorical}]
}

adreca_config = {
    "name": "adreca",
    "data": df_adreca,
    "tabular_model_configuration": {
        "max_training_time": 45
    },
    "primary_key": "ROWID_OBJECT_ADRECA",
    "columns": [{"name" : "ROWID_OBJECT_ADRECA", "model_encoding_type": ModelEncodingType.auto},
                {"name" : "IDENTIFICADOR", "model_encoding_type": ModelEncodingType.tabular_character}]
}

generator_config = {
    "name": "Multi-table Generator",
    "tables": [personas_table_config, documents_personals_table_config, persona_adreca_config, adreca_config] #contactes_config,
}

# COMMAND ----------

# MAGIC %md
# MAGIC ## Train

# COMMAND ----------

generator = mostly.train(name = "ctti-dSocial", config=generator_config, start=True, wait=False)

# COMMAND ----------

generator.training.wait(interval=60)

# COMMAND ----------

# MAGIC %md
# MAGIC ## Save Generator

# COMMAND ----------

generator.export_to_file('/Volumes/admin_govern_sta_des/tmp_mostlyai/generator/dSocial_03_04.zip')